# Copyright (c) OpenMMLab. All rights reserved.
import torch.nn.functional as F
from mmcv.cnn import ConvModule
from mmcv.cnn.bricks import NonLocal2d,ContextBlock
from mmcv.runner import BaseModule
from .model.attention.ECAAttention import ECAAttention
from ..builder import NECKS
from gsa_pytorch import GSA
from .model.attention.PSA import PSA
from .model.attention.EMSA import EMSA
from .model.attention.BAM import BAMBlock
from .model.attention.ShuffleAttention import ShuffleAttention
from .model.attention.SEAttention import SEAttention
from .model.attention.SKAttention import SKAttention
from .model.attention.CoTAttention import CoTAttention
from .model.attention.CBAM import CBAMBlock
from .model.attention.A2Atttention import DoubleAttention
@NECKS.register_module()
class GFP(BaseModule):
    """BFP (Balanced Feature Pyramids)

    BFP takes multi-level features as inputs and gather them into a single one,
    then refine the gathered feature and scatter the refined results to
    multi-level features. This module is used in Libra R-CNN (CVPR 2019), see
    the paper `Libra R-CNN: Towards Balanced Learning for Object Detection
    <https://arxiv.org/abs/1904.02701>`_ for details.

    Args:
        in_channels (int): Number of input channels (feature maps of all levels
            should have the same channels).
        num_levels (int): Number of input feature levels.
        conv_cfg (dict): The config dict for convolution layers.
        norm_cfg (dict): The config dict for normalization layers.
        refine_level (int): Index of integration and refine level of BSF in
            multi-level features from bottom to top.
        refine_type (str): Type of the refine op, currently support
            [None, 'conv', 'non_local'].
        init_cfg (dict or list[dict], optional): Initialization config dict.
    """

    def __init__(self,
                 in_channels,
                 num_levels,
                 refine_level=2,
                 refine_type=None,
                 conv_cfg=None,
                 norm_cfg=None,
                 init_cfg=dict(
                     type='Xavier', layer='Conv2d', distribution='uniform')):
        super(GFP, self).__init__(init_cfg)
        assert refine_type in [None, 'conv', 'non_local','gcblock','eca','gsa','bam','sa','se','sk','cot','cbam','a2']

        self.in_channels = in_channels
        self.num_levels = num_levels
        self.conv_cfg = conv_cfg
        self.norm_cfg = norm_cfg

        self.refine_level = refine_level
        self.refine_type = refine_type
        assert 0 <= self.refine_level < self.num_levels

        if self.refine_type == 'conv':
            self.refine = ConvModule(
                self.in_channels,
                self.in_channels,
                3,
                padding=1,
                conv_cfg=self.conv_cfg,
                norm_cfg=self.norm_cfg)
        elif self.refine_type == 'non_local':
            self.refine = NonLocal2d(
                self.in_channels,
                reduction=1,
                use_scale=False,
                conv_cfg=self.conv_cfg,
                norm_cfg=self.norm_cfg)
        elif self.refine_type == 'gcblock':
            self.refine = ContextBlock(
                self.in_channels,
                ratio=1
            )
        elif self.refine_type == 'eca':
            self.refine = ECAAttention(kernel_size=3)
        elif self.refine_type == 'gsa':
            self.refine = GSA(
            dim = self.in_channels,
            dim_out = self.in_channels,
            dim_key = 32,
            heads = 8,
            rel_pos_length = 256  # in paper, set to max(height, width). you can also turn this off by omitting this line
        )
        elif self.refine_type =='bam':
            self.refine = BAMBlock(channel=256,reduction=8,dia_val=1)
        elif self.refine_type =='sa':
            self.refine = ShuffleAttention(channel=256,G=8)
        elif self.refine_type =='se':
            self.refine = SEAttention(channel=256,reduction=8)
        elif self.refine_type == 'sk':
            self.refine = SKAttention(channel=256, reduction=8)
        elif self.refine_type =='cot':
            self.refine = CoTAttention(dim=256,kernel_size=3)
        elif self.refine_type =='cbam':
            self.refine = CBAMBlock(channel=256,reduction=16,kernel_size=5)
        elif self.refine_type =='a2':
            self.refine = DoubleAttention(256,64,64,True)
    def forward(self, inputs):
        """Forward function."""
        assert len(inputs) == self.num_levels

        # step 1: gather multi-level features by resize and average
        feats = []
        gather_size = inputs[self.refine_level].size()[2:]
        for i in range(self.num_levels):
            if i < self.refine_level:
                gathered = F.adaptive_max_pool2d(
                    inputs[i], output_size=gather_size)
            else:
                gathered = F.interpolate(
                    inputs[i], size=gather_size, mode='nearest')
            feats.append(gathered)

        bsf = sum(feats) / len(feats)

        # step 2: refine gathered features
        if self.refine_type is not None:
            bsf = self.refine(bsf)

        # step 3: scatter refined features to multi-levels by a residual path
        outs = []
        for i in range(self.num_levels):
            out_size = inputs[i].size()[2:]
            if i < self.refine_level:
                residual = F.interpolate(bsf, size=out_size, mode='nearest')
            else:
                residual = F.adaptive_max_pool2d(bsf, output_size=out_size)
            outs.append(residual + inputs[i])

        return tuple(outs)
